package com.hk.core.authentication.api.supports;

import com.hk.commons.util.CollectionUtils;
import com.hk.commons.util.IDGenerator;
import com.hk.core.authentication.api.AuthenticationException;
import com.hk.core.authentication.api.TokenExpiredException;
import com.hk.core.authentication.api.TokenGenerator;
import com.nimbusds.jose.JOSEException;
import com.nimbusds.jose.JWSAlgorithm;
import com.nimbusds.jose.JWSHeader;
import com.nimbusds.jose.crypto.RSASSASigner;
import com.nimbusds.jose.crypto.RSASSAVerifier;
import com.nimbusds.jose.jwk.RSAKey;
import com.nimbusds.jose.jwk.gen.RSAKeyGenerator;
import com.nimbusds.jwt.JWTClaimsSet;
import com.nimbusds.jwt.SignedJWT;
import lombok.SneakyThrows;

import java.text.ParseException;
import java.util.Collections;
import java.util.Date;
import java.util.Map;
import java.util.Objects;
import java.util.concurrent.TimeUnit;


/**
 * nimbus-jose-jwt
 */
public class SignedJWTTokenGenerator implements TokenGenerator {

    private static final RSAKey RSA_KEY;

    static {
        try {
            RSA_KEY = new RSAKeyGenerator(4096)
                    .generate();
        } catch (JOSEException e) {
            throw new RuntimeException(e);
        }
    }

    @Override
    @SneakyThrows
    public String generateToken(Map<String, Object> claims, long ttl, TimeUnit unit) {
        Date now = new Date();
        JWTClaimsSet.Builder builder = new JWTClaimsSet.Builder();
        builder.notBeforeTime(now)
                .jwtID(IDGenerator.uuid32().generate())
                .expirationTime(Date.from(now.toInstant().plus(ttl, unit.toChronoUnit())))
                .issueTime(now);
        if (CollectionUtils.isNotEmpty(claims)) {
            claims.forEach(builder::claim);
        }
        JWTClaimsSet jwtClaims = builder.build();
        RSASSASigner signer = new RSASSASigner(RSA_KEY.toRSAPrivateKey());
        SignedJWT jwt = new SignedJWT(new JWSHeader(JWSAlgorithm.RS512), jwtClaims);
        jwt.sign(signer);
        return jwt.serialize();
    }

    @Override
    @SneakyThrows
    public boolean verify(String token) {
        try {
            SignedJWT jwt = SignedJWT.parse(token);
            return jwt.verify(new RSASSAVerifier(RSA_KEY.toRSAPublicKey()));
        } catch (ParseException | JOSEException e) {
            return false;
        }
    }

    @Override
    public Map<String, Object> getClaims(String token) {
        try {
            SignedJWT jwt = SignedJWT.parse(token);
            JWTClaimsSet jwtClaims = jwt.getJWTClaimsSet();
            if (Objects.isNull(jwtClaims)) {
                return Collections.emptyMap();
            }
            Date expirationTime = jwtClaims.getExpirationTime();
            if (Objects.nonNull(expirationTime) && new Date().after(expirationTime)) {
                throw new TokenExpiredException(token, expirationTime.toInstant());
            }
            return jwtClaims.getClaims();
        } catch (ParseException e) {
            throw new AuthenticationException("error token", e);
        }
    }
}
